/*
 * Copyright 2012 The LoadTestAnalyzer Project
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 */

package org.yantsu.loadtestanalyzer.ltaweb;

import static org.elasticsearch.node.NodeBuilder.nodeBuilder;
import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
import static org.elasticsearch.index.query.QueryBuilders.*;
import static org.elasticsearch.index.query.FilterBuilders.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.apache.log4j.Logger;
import org.elasticsearch.action.admin.indices.optimize.OptimizeRequest;
import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.action.support.replication.ReplicationType;
import org.elasticsearch.client.Client;
import org.elasticsearch.client.Requests;
import org.elasticsearch.cluster.metadata.IndexMetaData;
import org.elasticsearch.cluster.metadata.MappingMetaData;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.node.Node;
import org.elasticsearch.node.NodeBuilder;

public class ElasticSearchClient {
	private static final String DATA_PATH_PROP = "de.yantsu.lta.datapath";
	private static final String LTA_INDEX_NAME = "ltadata";

	private static Logger LOG = Logger.getLogger(ElasticSearchClient.class);
	
	private Node node;
	private Client client;

	public ElasticSearchClient() {
		String dataPath = System.getProperty(DATA_PATH_PROP);
		NodeBuilder nodeBuilder = nodeBuilder();
		if(dataPath != null) {
			nodeBuilder.settings().put("path.data", dataPath);
		}

		nodeBuilder.settings().put("index.cache.field.type", "soft");
		
		node = nodeBuilder.node();
		client = node.client();			
		client.admin().cluster().health(Requests.clusterHealthRequest().waitForYellowStatus()).actionGet();	
		if(!client.admin().indices().exists(Requests.indicesExistsRequest(LTA_INDEX_NAME)).actionGet().isExists()) {
			client.admin().indices().create(Requests.createIndexRequest(LTA_INDEX_NAME)
					.mapping("series", "{\"series\": {\"_id\": {\"index\": \"not_analyzed\"}}}")).actionGet();
			client.admin().indices().putMapping(Requests.putMappingRequest(LTA_INDEX_NAME).type("data_Grinder").source(
					"{\"data_Grinder\" : {" +
							"\"_parent\" : {\"type\" : \"series\"}," +
							"\"properties\" : {" +
								"\"timestamp\" : {\"type\" : \"date\"}," +
								"\"category\" : {\"type\" : \"string\", \"index\" : \"not_analyzed\"}," +
								"\"hostname\" : {\"type\" : \"string\", \"index\" : \"not_analyzed\"}" +
							"}" +
						"}" +
					"}")).actionGet();
			client.admin().indices().putMapping(Requests.putMappingRequest(LTA_INDEX_NAME).type("data_Sar").source(
					"{\"data_Sar\" : {" +
							"\"_parent\" : {\"type\" : \"series\"}," +
							"\"properties\" : {" +
								"\"timestamp\" : {\"type\" : \"date\"}," +
								"\"category\" : {\"type\" : \"string\", \"index\" : \"not_analyzed\"}," +
								"\"hostname\" : {\"type\" : \"string\", \"index\" : \"not_analyzed\"}" +
							"}" +
						"}" +
					"}")).actionGet();
			client.admin().cluster().health(Requests.clusterHealthRequest(LTA_INDEX_NAME).waitForYellowStatus()).actionGet();
		}
		if(LOG.isInfoEnabled()) {
			LOG.info("ElasticSearch Cluster ready.");
		}		
	}
	
	public String getData(String query) {
		return client.prepareSearch(LTA_INDEX_NAME).setSource(query).execute().actionGet().toString();
	}	
	
	@SuppressWarnings("unchecked")
	public Map<String, List<String>> getTypes() {
		IndexMetaData metaData = client.admin().cluster().state(Requests.clusterStateRequest()
				.filterRoutingTable(true)
				.filterNodes(true)
				.filteredIndices(LTA_INDEX_NAME)).actionGet().getState().metaData().iterator().next();
		Map<String, List<String>> types = new HashMap<String, List<String>>();
		for (MappingMetaData mapping : metaData.mappings().values()) {
			if(mapping.type().startsWith("data_")) {
				try {
					Map<String, Object> map = mapping.sourceAsMap();
					List<String> labels = new ArrayList<String>();
					try {
						Map<String, Object> properties = (Map<String, Object>)map.get("properties");
						Map<String, Object> values = (Map<String, Object>)properties.get("values");
						properties = (Map<String, Object>)values.get("properties");
						for (Iterator<String> iterator = properties.keySet().iterator(); iterator.hasNext();) {
							labels.add(iterator.next());
						}
					} catch(ClassCastException e) {
						throw new IllegalStateException("Reading types failed", e);
					}
					types.put(mapping.type(), labels);
				} catch (IOException e) {
					throw new IllegalStateException("Reading types failed", e);
				}
			}
		}
		return types;
	}
	
	public void deleteSeries(String pattern) {
		if(pattern != null && pattern.length() > 0) {
			client.prepareDeleteByQuery(LTA_INDEX_NAME).setQuery(
					constantScoreQuery(
						orFilter(
							andFilter(
								typeFilter("series"),
								queryFilter(wildcardQuery("_id", "*" + pattern + "*"))
							),
							queryFilter(wildcardQuery("_parent", "*" + pattern + "*"))
						)
					)
				).execute().actionGet();
		} else {
			client.prepareDeleteByQuery(LTA_INDEX_NAME).setQuery(matchAllQuery()).execute().actionGet();
		}
		client.admin().indices().refresh(new RefreshRequest(LTA_INDEX_NAME)).actionGet();
		client.admin().indices().optimize(new OptimizeRequest(LTA_INDEX_NAME)).actionGet();
	}
	
	public BulkRequestBuilder prepareBulk() {
		return client.prepareBulk();
	}

	/**
	 * 
	 * @param testname
	 * @param agenttype
	 * @return doc id
	 */
	public String getSeriesDocAndLock(String agenttype, String testname) {
		GetResponse response = client.get(Requests.getRequest(LTA_INDEX_NAME).id(testname)).actionGet();
		if(response.isExists()) {
			Object lock = response.getSource().get(agenttype + "_lock");
			boolean locked = lock != null && Boolean.parseBoolean(lock.toString());
			if(!locked) {
				response.getSource().put(agenttype + "_lock", true);
				client.prepareIndex(LTA_INDEX_NAME, "series", testname)
						.setVersion(response.getVersion())
						.setSource(response.getSource()).execute().actionGet();
			} else {
				throw new ConcurrentImportException("An import for the agenttype " + agenttype + " is already runnning");
			}
		} else {
			XContentBuilder builder;
			try {
				builder = jsonBuilder();
				builder.startObject();
					builder.field(agenttype + "_lock", true);
				builder.endObject();
				builder.close();
			} catch (IOException e) {
				throw new IllegalStateException("Error writing json doc", e);
			}
			client.prepareIndex(LTA_INDEX_NAME, "series")
					.setSource(builder)
					.setCreate(true)
					.setId(testname).execute().actionGet();
		}		
		return testname;
	}
	
	public void finalizeSeriesDoc(String testname, String agenttype, long start, long end) {
		GetResponse response = client.get(Requests.getRequest(LTA_INDEX_NAME).id(testname)).actionGet();
		if(response.isExists()) {
			response.getSource().put(agenttype + "_lock", false);
						
 			//calculate new overall start timestamp
 			Object startObj = response.getSource().get("start");
 			long oldStart = Long.MAX_VALUE;
 			if(startObj != null) {
 				try {
 					oldStart = Long.parseLong(startObj.toString());
 					if(oldStart > start ) {
 						response.getSource().put("start", start);
 					}
 				} catch(NumberFormatException e) {
 					throw new IllegalStateException("Illegal Format of start param - cannot finalize series data", e);
 				}
 			} else {
 				response.getSource().put("start", start);
 			}
 			
 			//calculate new overall end timestamp
 			Object endObj = response.getSource().get("end");
 			long oldEnd = Long.MIN_VALUE;
 			if(endObj != null) {
 				try {
 					oldEnd = Long.parseLong(endObj.toString());
 					if(oldEnd < end ) {
 						response.getSource().put("end", end);
 					}
 				} catch(NumberFormatException e) {
 					throw new IllegalStateException("Illegal Format of end param - cannot finalize series data", e);
 				}
 			} else {
 				response.getSource().put("end", end);
 			}

 			client.prepareIndex(LTA_INDEX_NAME, "series", testname)
					.setVersion(response.getVersion())
					.setSource(response.getSource()).execute().actionGet();
		} else {
			throw new IllegalStateException("Cannot finalize series data - doc not found");
		}
	}
	
	public IndexRequestBuilder prepareData(String agenttype, String id, XContentBuilder builder) {
		return client.prepareIndex(LTA_INDEX_NAME, "data_" + agenttype).setParent(id).setSource(builder);
	}
	
	public void writeBulk(BulkRequestBuilder bulkRequest) {
		bulkRequest.setReplicationType(ReplicationType.ASYNC);
		BulkResponse bulkResponse = bulkRequest.execute().actionGet();
		if (bulkResponse.hasFailures()) {
			throw new IllegalStateException("Writing data failed: " + bulkResponse.buildFailureMessage());
		}
	}	
}
